package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

// HttpsqsStatus 定义 httpsqs 返回的队列状态的结构体
type HttpsqsStatus struct {
	Name     string `json:"name"`
	Maxqueue int64  `json:"maxqueue"`
	Putpos   int64  `json:"putpos"`
	Putlap   int64  `json:"putlap"`
	Getpos   int64  `json:"getpos"`
	Getlap   int64  `json:"getlap"`
	Unread   int64  `json:"unread"`
}

// Exporter 是 Prometheus 采集器结构，保存监控信息和配置
type Exporter struct {
	httpsqsInstances []string // httpsqs服务地址列表，例如 192.168.4.117:9001
	queueNames       []string // 需要监控的队列名称列表

	mutex sync.Mutex // 访问指标时保护数据安全

	maxqueue *prometheus.GaugeVec
	putpos   *prometheus.GaugeVec
	putlap   *prometheus.GaugeVec
	getpos   *prometheus.GaugeVec
	getlap   *prometheus.GaugeVec
	unread   *prometheus.GaugeVec
}

// NewExporter 创建 Exporter 实例并初始化指标
func NewExporter(instances []string, queues []string) *Exporter {
	return &Exporter{
		httpsqsInstances: instances,
		queueNames:       queues,
		maxqueue: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Name: "httpsqs_maxqueue",
				Help: "Maximum capacity of the queue.",
			},
			[]string{"name", "instance"},
		),
		putpos: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Name: "httpsqs_putpos",
				Help: "Current write position index in the queue.",
			},
			[]string{"name", "instance"},
		),
		putlap: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Name: "httpsqs_putlap",
				Help: "Number of times the write pointer has looped over the queue.",
			},
			[]string{"name", "instance"},
		),
		getpos: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Name: "httpsqs_getpos",
				Help: "Current read position index in the queue.",
			},
			[]string{"name", "instance"},
		),
		getlap: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Name: "httpsqs_getlap",
				Help: "Number of times the read pointer has looped over the queue.",
			},
			[]string{"name", "instance"},
		),
		unread: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Name: "httpsqs_unread",
				Help: "Number of unread (unconsumed) messages in the queue.",
			},
			[]string{"name", "instance"},
		),
	}
}

// Describe 实现 prometheus.Collector 接口，用于描述指标
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	e.maxqueue.Describe(ch)
	e.putpos.Describe(ch)
	e.putlap.Describe(ch)
	e.getpos.Describe(ch)
	e.getlap.Describe(ch)
	e.unread.Describe(ch)
}

// fetchStatus 通过 httpsqs HTTP 接口拉取指定实例和队列的状态数据
func (e *Exporter) fetchStatus(instance, queue string) (*HttpsqsStatus, error) {
	url := fmt.Sprintf("http://%s?name=%s&opt=status_json", instance, queue)
	resp, err := http.Get(url)
	if err != nil {
		return nil, fmt.Errorf("http get error: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("bad status code: %d", resp.StatusCode)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("read body error: %v", err)
	}

	var status HttpsqsStatus
	if err := json.Unmarshal(body, &status); err != nil {
		return nil, fmt.Errorf("json unmarshal error: %v", err)
	}

	return &status, nil
}

// Collect 实现 prometheus.Collector 接口，用于采集指标数据
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	e.mutex.Lock()
	defer e.mutex.Unlock()

	for _, instance := range e.httpsqsInstances {
		for _, queue := range e.queueNames {
			status, err := e.fetchStatus(instance, queue)
			if err != nil {
				log.Printf("Error fetching status from %s queue %s: %v", instance, queue, err)
				continue
			}

			labels := prometheus.Labels{"name": status.Name, "instance": instance}
			e.maxqueue.With(labels).Set(float64(status.Maxqueue))
			e.putpos.With(labels).Set(float64(status.Putpos))
			e.putlap.With(labels).Set(float64(status.Putlap))
			e.getpos.With(labels).Set(float64(status.Getpos))
			e.getlap.With(labels).Set(float64(status.Getlap))
			e.unread.With(labels).Set(float64(status.Unread))
		}
	}

	e.maxqueue.Collect(ch)
	e.putpos.Collect(ch)
	e.putlap.Collect(ch)
	e.getpos.Collect(ch)
	e.getlap.Collect(ch)
	e.unread.Collect(ch)
}

// splitAndTrim 以逗号分割字符串并去除空白
func splitAndTrim(s string) []string {
	parts := strings.Split(s, ",")
	var ret []string
	for _, p := range parts {
		p = strings.TrimSpace(p)
		if p != "" {
			ret = append(ret, p)
		}
	}
	return ret
}

func main() {
	listenAddr := flag.String("listen", ":9101", "Address to listen on for /metrics")
	instancesStr := flag.String("instances", "", "Comma-separated httpsqs service addresses (required)")
	queueNamesStr := flag.String("queues", "", "Comma-separated queue names to monitor (required)")
	flag.Parse()

	if *instancesStr == "" {
		fmt.Fprintln(os.Stderr, "Error: -instances parameter is required")
		flag.Usage()
		os.Exit(1)
	}
	if *queueNamesStr == "" {
		fmt.Fprintln(os.Stderr, "Error: -queues parameter is required")
		flag.Usage()
		os.Exit(1)
	}

	instanceList := splitAndTrim(*instancesStr)
	queueList := splitAndTrim(*queueNamesStr)

	exporter := NewExporter(instanceList, queueList)
	prometheus.MustRegister(exporter)

	http.Handle("/metrics", promhttp.Handler())
	log.Printf("Starting httpsqs exporter on %s", *listenAddr)
	log.Fatal(http.ListenAndServe(*listenAddr, nil))
}
